# !/usr/bin/env python
# -*- coding: utf-8 -*-
# @File  : flask_rabbitmq实现异步任务功能.py
# @Author: dongguangwen
# @Date  : 2024-12-28 21:16
import json
import threading
import time
from flask import Flask, request, jsonify
import pika

app = Flask(__name__)

# RabbitMQ配置
RABBITMQ_HOST = '192.168.1.119'
RABBITMQ_USER = 'admin'  # 修改为你的用户名
RABBITMQ_PASSWORD = 'password'  # 修改为你的密码
QUEUE_NAME = 'async_queue'


# RabbitMQ连接
def get_rabbitmq_connection():
    credentials = pika.PlainCredentials(RABBITMQ_USER, RABBITMQ_PASSWORD)
    return pika.BlockingConnection(pika.ConnectionParameters(host=RABBITMQ_HOST, credentials=credentials))


# 后台异步任务处理
def process_message(body):
    task_data = json.loads(body)
    print(f"[Worker] Processing task: {task_data}")
    time.sleep(5)  # 模拟任务耗时
    print(f"[Worker] Task completed: {task_data}")


def start_worker():
    """后台工作线程"""
    def worker():
        connection = get_rabbitmq_connection()
        channel = connection.channel()
        channel.queue_declare(queue=QUEUE_NAME)

        def callback(ch, method, properties, body):
            print(f"[Worker] Received: {body}")
            process_message(body)
            ch.basic_ack(delivery_tag=method.delivery_tag)

        channel.basic_qos(prefetch_count=1)  # 确保公平分发任务
        channel.basic_consume(queue=QUEUE_NAME, on_message_callback=callback)

        print("[Worker] Waiting for messages...")
        channel.start_consuming()

    thread = threading.Thread(target=worker, daemon=True)
    thread.start()


# API路由
@app.route('/submit_task', methods=['POST'])
def submit_task():
    """提交异步任务"""
    try:
        task_data = request.json
        if not task_data:
            return jsonify({'status': 'error', 'message': 'Invalid task data.'}), 400

        connection = get_rabbitmq_connection()
        channel = connection.channel()
        channel.queue_declare(queue=QUEUE_NAME)

        message = json.dumps(task_data)
        channel.basic_publish(exchange='', routing_key=QUEUE_NAME, body=message)

        connection.close()
        return jsonify({'status': 'success', 'message': 'Task submitted successfully.'}), 200
    except Exception as e:
        return jsonify({'status': 'error', 'message': str(e)}), 500


@app.route('/status', methods=['GET'])
def status():
    """检查服务状态"""
    return jsonify({'status': 'running', 'message': 'Service is up and running!'}), 200


if __name__ == '__main__':
    start_worker()  # 启动后台消费线程
    app.run()
